#!/usr/bin/env python
# encoding: utf-8

import networkx as nx
import matplotlib.pyplot as plt

import sys
import os
import random
import math

from protocols import Simple
from protocols import Crowds
from protocols import OnionRouting
from protocols import Utilities


"""
'
' Author: Jonathan Potter
'
' Name: Driver
'
' Description: Drives the simulation, including reading in simulation parameters,
' generating a network, running each protocol on the generated network, and
' graphing statistics.
'
"""
class Driver:
	
	# The simulation input parameters (knobs)
	parameters = None
	
	# The list of protocols to run the simulation with
	protocols = None
	
	"""
	' Name: __init__
	'
	' Description: Constructor for class Driver.  Sets up protocols and parameters.
	' 		
	' Parameters: parameters The simulation wide parameters (knobs)
	'
	' Return: An instance of class Driver.
	'
	"""
	def __init__(self, parameters):
		self.parameters = parameters
		self.protocols = {
			'Simple': Simple(),
			'Crowds': Crowds(),
			'OnionRouting': OnionRouting()
		}		

	"""
	' Name: generateNetwork
	'
	' Description: Generates a random network to use in the simulation.
	' 		
	' Parameters: numberOfNodes The number of nodes, or network size of the generated network.
	'
	' Return: A network of size numberOfNodes 
	'
	"""	
	def generateNetwork(self, numberOfNodes):
		network = nx.connected_watts_strogatz_graph(numberOfNodes, int(round(math.log(numberOfNodes))), .8)
		
		# add a random weight between 1 and 10 for each edge
		for (u,v) in network.edges():
			weight = random.randint(*self.parameters['bandwidthRange'])
			network[u][v] = {'weight': weight}
			network[v][u] = {'weight': weight}
		
		return network
	
	
	"""
	' Name: displayNetwork
	'
	' Description: Saves a visual representation of the network to a file.  This includes a highlighted path.
	' 		
	' Parameters:	network The network to display.
	'				path The path to highlight on the network.
	'				filename The filename to save the image to.
	'
	' Return: None.
	'
	"""
	def displayNetwork(self, network, path, filename):
		(nodeColors, edgeColors) = self.highlightPath(network, path)
		
		# Do widths
		edgeWidths = []
		for (u, v) in network.edges():
			edgeWidths.append(Utilities.getBandwidth(self.parameters, network, u, v))
		
		pos = nx.graphviz_layout(network, prog = 'dot')
				
		nx.draw_networkx_edges(	network,
								pos,
								alpha = .3,
								width = edgeWidths,
								edge_color = edgeColors)
								
		nx.draw_networkx_edges(	network,
								pos,
								alpha = .4,
								width = 1,
								edge_color = 'k')
		
		nx.draw_networkx_nodes(	network, 
								pos, 
								node_size = 80, 
								alpha = 1, 
								node_color = nodeColors, 
								with_labels = False)
		
		plt.axis('off')
		
		# draw it
		plt.savefig(filename)
		plt.close()
	
	
	
	"""
	' Name: getRandomEndpoints
	'
	' Description: Gets a random source and destination node in the network to use for routing data
	' in the simulation.
	' 		
	' Parameters:	network The network to get source and destination from.
	'
	' Return: (u, v) U is the source node, and V is the destination node.
	'
	"""	
	def getRandomEndpoints(self, network):
		u = random.choice(network.nodes())
		
		withoutU = network.nodes()[:]
		withoutU.remove(u)
		
		v = random.choice(withoutU)
		
		return (u, v)
	
	"""
	' Name: highlightPath
	'
	' Description: Visually highlights a path in our network.  This function is used when saving
	' a snapshot to a file.
	' 		
	' Parameters:	network The network to use.
	'				path The path to highlight on the network.
	'
	' Return: (node_colors, edge_colors) Where node_colors and edge_colors are lists of color information to
	' be used when saving a snapshot of the graph.  These have the path highlighted.
	'
	"""
	def highlightPath(self, network, path):
		node_colors = [255 for i in range(network.number_of_nodes())]
		edge_colors = ['black' for i in range(network.number_of_edges())]
		
		pairs = Utilities.pairs(path)[:-1]
		
		for (a, b) in pairs:
			if b < a:
				edge = (b, a)
			else:
				edge = (a, b)
			
			try:
				aIndex = network.nodes().index(a)
				bIndex = network.nodes().index(b)
				node_colors[aIndex] = 0
				node_colors[bIndex] = 0
			except:
				pass
					
			try:
				edgeIndex = network.edges().index(edge)
				edge_colors[edgeIndex] = 'magenta'
			except:
				pass
		
		return (node_colors, edge_colors)
	
	"""
	' Name: initPlot
	'
	' Description: Initialize gnuplot to plot a certain dataset.  This function opens a pipe to gnuplot and
	' sets up basic axes and titles, etc.
	' 		
	' Parameters:	name The name of this plot.  This will be used as the filename to store the raw data in.
	'				title The title of this plot.  This is used as the title of the graph.
	'				xLabel The label for the x axis of the plot.
	'				yLabel The label for the y axis of the plot.
	'
	' Return: (plotFile, gnuplot) Handles to the data file, and the gnuplot pipe respectively.
	'
	"""
	def initPlot(self, name, title, xLabel, yLabel):
		# initialize plot file	
		plotFile = open(name, "w")
				
		gnuplot = os.popen("gnuplot -persist 2> gnuplot.log", "w")
		
		# initialize gnuplot
		# gnuplot.write("set terminal x11\n")
		gnuplot.write("set title '%s'\n" % title)
		gnuplot.write("set xlabel '%s'\n" % xLabel)
		gnuplot.write("set ylabel '%s'\n" % yLabel)
		gnuplot.write("plot '%s' using 1:2 title 'Simple' with lines, '%s' using 1:3 title 'Crowds' with lines, '%s' using 1:4 title 'Onion Routing' with lines\n" % (name, name, name))
		
		return (plotFile, gnuplot)
	
	"""
	' Name: plotSlice
	'
	' Description: Plots one datapoint of the graph.  This function is called each time through a loop
	' to get live real time plotting.
	' 		
	' Parameters:	plot The plot returned by the initPlot function.
	'				data The datapoint to plot.
	'
	' Return: None.
	'
	"""
	def plotSlice(self, plot, data):
		(plotFile, gnuplot) = plot
		
		plotFile.write(' '.join(map(lambda x: str(x), data)) + "\n")
		plotFile.flush()
		
		gnuplot.write("replot\n")
		try:
			gnuplot.flush()
		except:
			pass

	
		"""
		' Name: closePlot
		'
		' Description: Cleans up all plot variables, and closes file and pipe handles.
		' 		
		' Parameters:	plot The plot returned by initPlot.
		'
		' Return: None.
		'
		"""
	def closePlot(self, plot):
		(plotFile, gnuplot) = plot

		plotFile.close()
		gnuplot.close()
	
	
	"""
	' Name: simulate
	'
	' Description: Main simulation funciton.  Drives each of the protocols, generates the networks, and plots
	' statistics from the simulation.
	' 		
	' Parameters: None
	'
	' Return: None.
	'
	"""
	def simulate(self):
		bandwidthPlot = self.initPlot("bandwidth", "Bandwidth Comparison", "Number of nodes", "Bandwidth (Mbps)")
		overheadPlot = self.initPlot("overhead", "Overhead Comparison", "Number of nodes", "Overhead (milliseconds)")
		
		# data in bits
		amountOfData = self.parameters['defaultDataNumberOfPackets'] * self.parameters['packetSize'] * 8
		
		for n in range(*self.parameters['numberOfNodesRange']):
			iterationTimes = [[], [], []]
			
			for i in range(self.parameters['iterations']):
				network = self.generateNetwork(n)
				(source, destination) = self.getRandomEndpoints(network)
				
				# Simple protocol
				simplePath = self.protocols['Simple'].getPath(network, self.parameters, source, destination)
				iterationTimes[0].append(self.protocols['Simple'].runSimulation(network, self.parameters, simplePath))
				
				# Crowds protocol
				(crowdsPath, shortestPath, crowdsNetwork) = self.protocols['Crowds'].getPath(network, self.parameters, source)
				iterationTimes[1].append(self.protocols['Crowds'].runSimulation(crowdsNetwork, self.parameters, crowdsPath, shortestPath))
				self.protocols['Crowds'].merge(crowdsPath, shortestPath)

				# Onion Routing protocol
				
				orPath = self.protocols['OnionRouting'].getPath(self.parameters, network, source, destination)
				iterationTimes[2].append(self.protocols['OnionRouting'].runSimulation(network, self.parameters, orPath))

				if n in self.parameters['snapshot'] and i is 1:
					self.displayNetwork(network, simplePath, 'nodes_%s_protocol_%s' % (n, 'simple'))
					self.displayNetwork(crowdsNetwork, crowdsPath, 'nodes_%s_protocol_%s' % (n, 'crowds'))
					self.displayNetwork(network, orPath, 'nodes_%s_protocol_%s' % (n, 'onionrouting'))
			
		
			# In ms
			averageTimes = map(lambda p: float(sum(p)) / len(p), iterationTimes)
			
			# In Mbps
			bandwidths = map(lambda x: (float(amountOfData) / x) / 10 ** 3, averageTimes)
			
			# In ms
			overheads = averageTimes
			
			self.plotSlice(bandwidthPlot, [n] + bandwidths)
			self.plotSlice(overheadPlot, [n] + overheads)
		
		self.closePlot(bandwidthPlot)
		self.closePlot(overheadPlot)
	


